import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter.ttk import Progressbar
from PIL import Image, ImageTk
import face_recognition
import os
import threading
import queue
import time
import csv
import openpyxl
from openpyxl.styles import Alignment # 用于导出Excel
import shutil  # 用于复制文件
import logging
import psutil  # 用于监控内存和CPU使用情况
from datetime import datetime
import packaging
import packaging.version
import packaging.specifiers
import packaging.requirements

# 目前存在问题：继续按钮是跳过了已经在目录里的内容，但是条清空了，并且好像进程有问题
# 待完善功能：操作日志和报错日志（已解决） 导出目录（已解决）  生成单独的软件（已解决）
# 已添加功能：完善基本UI设计：调整展示图片区域，添加列表框滚动条，优化预览图片速度，增加执行按钮
#           优化搜索流程：展示选定搜索目录与实时搜索目录，修复搜索图片数量不精准的问题
#           提供导出CSV，EXCEL格式的搜索记录，提供导出已匹配的图片文件夹
#           保存详细的操作记录日志：用户操作记录、性能监控、匹配结果记录、错误记录
#           提供图片相似度展示
#

# 创建主窗口
root = tk.Tk()
root.title("人脸识别图片搜索工具")
root.geometry("1200x600")

# 全局变量
target_image_path = None
target_encoding = None
matched_images = set()
matched_images_list = []  # 列表来存储匹配结果用于显示
matched_images_with_similarity_list = []  # 用于导出路径和相似度
error_images = []
selected_directory = tk.StringVar()  # 用于保存选择的目录
current_directory = tk.StringVar() # 用于保存当前正在扫描的目录路径
is_searching = False  # 用于控制搜索线程的终止
progress_queue = queue.Queue()  # 用于线程间通信
stop_event = threading.Event()
processed_files = []  # 保存已经处理的文件路径
start_time = None # 在搜索开始时记录开始时间

# 配置日志记录
logging.basicConfig(
    filename="face_recognition_search.log",  # 日志文件名
    level=logging.INFO,  # 记录信息级别为INFO及以上
    format="%(asctime)s - %(levelname)s - %(message)s",  # 日志格式
    datefmt="%Y-%m-%d %H:%M:%S"  # 日期格式
)

# 记录程序启动
logging.info('--------------------------------------------------------------------------------------------------')
logging.info("用户打开程序")

# 日志记录函数
def log_message(message, level="info"):
    if level == "info":
        logging.info(message)
    elif level == "warning":
        logging.warning(message)
    elif level == "error":
        logging.error(message)

# 选择目标图片的函数
def choose_target_image():
    global target_image_path, target_encoding
    target_image_path = filedialog.askopenfilename(
        title="选择目标人脸图片",
        filetypes=[("Image Files", "*.jpg;*.jpeg;*.png")]
    )
    if target_image_path:
        logging.info(f"用户选择目标图片路径: {target_image_path}")
        load_target_image(target_image_path)
        try:
            target_image = face_recognition.load_image_file(target_image_path)
            encodings = face_recognition.face_encodings(target_image)
            if encodings:
                target_encoding = encodings[0]
            else:
                messagebox.showerror("错误", "目标图片中没有检测到人脸。")
                target_image_path = None
                target_encoding = None
                target_img_label.config(image=None)
        except Exception as e:
            messagebox.showerror("错误", f"加载图片失败：{e}")

# 加载目标图片并显示
def load_target_image(image_path):
    try:
        img = Image.open(image_path)
        img = img.resize((200, 200), Image.ANTIALIAS)
        img = ImageTk.PhotoImage(img)
        target_img_label.config(image=img)
        target_img_label.image = img
    except Exception as e:
        messagebox.showerror("错误", f"加载目标图片失败：{e}")

# 选择要搜索的文件夹
def choose_search_directory():
    global processed_files  # 需要更新全局变量
    if not target_image_path:
        messagebox.showerror("错误", "请先选择目标图片。")
        return

    search_directory = filedialog.askdirectory(title="选择搜索目录")
    if search_directory:
        logging.info(f"用户选择了搜索目录: {search_directory}")
        selected_directory.set(search_directory)  # 更新标签内容
        disable_buttons()
        progress_bar['value'] = 0  # 重置进度条
        # 启动线程进行搜索
        search_thread = threading.Thread(target=search_for_matches, args=(search_directory,))
        search_thread.start()

# 搜索匹配图片
def search_for_matches(search_directory):
    global matched_images, error_images, matched_images_list, is_searching, start_time
    matched_images.clear()  # 清空集合
    error_images.clear()
    matched_images_list.clear()  # 清空列表
    stop_event.clear()  # 清除停止事件
    result_list.delete(0, tk.END)
    error_images = []
    is_searching = True

    log_message(f"搜索开始")

    supported_formats = (".jpg", ".jpeg", ".png",".JPG",".JPEG",".PNG")
    total_images = 0

    start_time = time.time()  # 记录搜索开始时间

    # 记录用户点击了开始搜索按钮
    logging.info("用户点击开始搜索按钮")

    # 记录开始搜索时间和系统状态
    memory_info = psutil.virtual_memory()
    cpu_usage = psutil.cpu_percent()
    logging.info(f"搜索开始，当前内存使用: {memory_info.used // (1024 * 1024)} MB, CPU 使用率: {cpu_usage}%")

    logging.info(f"搜索目录: {search_directory}, 是否递归搜索子目录: 是")  # 假设递归搜索

    # 统计总的支持的图片数量
    for root_dir, _, files in os.walk(search_directory):
        if stop_event.is_set():  # 检查是否应停止
            break
        for file in files:
            if file.endswith(supported_formats):
                total_images += 1  # 计算支持格式的文件数量

    log_message(f"总图片数量: {total_images}")

    if total_images == 0:
        log_message("在选定的目录中没有支持的图片格式。")
        messagebox.showinfo("结果", "在选定的目录中没有支持的图片格式。")
        enable_buttons()
        return

    print(f"当前搜索目录为 {search_directory}") # 打印搜索目录
    print(f"总共有 {total_images} 张支持的图片。")  # 打印总数

    processed_count = 0

    for root_dir, _, files in os.walk(search_directory):
        if stop_event.is_set():  # 检查是否应停止
            break
        current_directory.set(root_dir)  # 更新当前正在扫描的目录
        for file in files:
            if stop_event.is_set():  # 检查是否应停止
                break
            image_path = os.path.join(root_dir, file)
            if image_path in processed_files:  # 如果已经处理过，跳过
                continue
            if file.endswith(supported_formats):
                try:
                    # 检查图像大小
                    img = Image.open(image_path)
                    img.verify()  # 验证图像是否损坏
                    img.close()

                    # 处理图像
                    image = face_recognition.load_image_file(image_path)
                    encodings = face_recognition.face_encodings(image)

                    if encodings:  # 只有在找到人脸时才进行匹配
                        for encoding in encodings:
                            face_distance = face_recognition.face_distance([target_encoding], encoding)
                            similarity = (1 - face_distance[0]) * 100  # 转换为相似度百分比
                            if face_distance[0] < 0.6:  # 使用距离判断相似性
                                matched_images.add(image_path)  # 添加到集合中
                                matched_images_list.append(f"{image_path} ")# 用于右侧展示的图片链接
                                matched_images_with_similarity_list.append(
                                    (image_path, f"{similarity:.2f}%"))  # 保存路径和相似度用于导出
                                result_list.insert(tk.END, f"{image_path}   [相似度: {similarity:.2f}%]") # 在列表框中显示

                    processed_files.append(image_path)  # 添加到已处理文件列表
                    processed_count += 1
                    progress_queue.put({'processed_count': processed_count, 'total_images': total_images})

                except Exception as e:
                    error_images.append(image_path)  # 记录处理失败的图像路径
                    log_message(f"处理图片失败: {image_path}, 错误: {str(e)}", level="error")

    total_matched_images = len(matched_images_with_similarity_list)  # 匹配图片总数
    logging.info(f"搜索完成，扫描图片总数: {total_images} 张，匹配图片总数: {total_matched_images} 张")

    on_search_complete()

# 更新进度条和标签
def update_ui_from_queue():
    try:
        while True:
            progress_info = progress_queue.get_nowait()
            progress_var.set(f"已处理 {progress_info['processed_count']}/{progress_info['total_images']} 张图片")
            progress_bar['value'] = (progress_info['processed_count'] / progress_info['total_images']) * 100
    except queue.Empty:
        pass
    root.after(1000, update_ui_from_queue)  # 100ms 后再次检查队列

# 搜索完成后的处理逻辑
def on_search_complete():

    global start_time
    end_time = time.time()  # 搜索结束时间
    elapsed_time = end_time - start_time  # 计算搜索时间

    # 记录搜索结束
    elapsed_time = round(time.time() - start_time, 2)
    final_memory_info = psutil.virtual_memory()
    final_cpu_usage = psutil.cpu_percent()
    logging.info(f"搜索结束，用时: {elapsed_time}秒，最终内存使用: {final_memory_info.used // (1024 * 1024)} MB, CPU 使用率: {final_cpu_usage}%")


    if not matched_images:
        messagebox.showinfo("结果", "没有匹配的图片。")
    else:
        messagebox.showinfo("结果", f"找到 {len(matched_images)} 张匹配的图片。")

    if error_images:
        messagebox.showwarning("处理警告", f"有 {len(error_images)} 张图片处理失败。")

    enable_buttons()

# 添加一个查看日志的按钮
def open_log_file():
    try:
        os.startfile("face_recognition_search.log")  # 在Windows中打开日志文件
    except Exception as e:
        messagebox.showerror("错误", f"无法打开日志文件: {e}")

# 导出所有匹配到的图片
def export_matched_images():
    if not matched_images_list:
        messagebox.showinfo("提示", "没有可以导出的匹配结果。")
        return

    export_dir = filedialog.askdirectory(title="选择导出文件夹")
    if not export_dir:
        return

    try:
        for matched_image in matched_images_list:
            # 获取文件名
            file_name = os.path.basename(matched_image)
            # 生成目标路径
            target_path = os.path.join(export_dir, file_name)
            # 复制文件
            shutil.copy(matched_image, target_path)

        messagebox.showinfo("成功", f"所有匹配的图片已导出到 {export_dir}")
        logging.info(f"匹配图片文件夹已导出至 {export_dir}, 文件大小: {os.path.getsize(export_dir) // 1024} KB")

    except Exception as e:
        messagebox.showerror("错误", f"导出失败: {e}")

# 导出CSV
def export_to_csv():
    if not matched_images_list:
        messagebox.showinfo("提示", "没有可以导出的匹配结果。")
        return

    export_path = filedialog.asksaveasfilename(defaultextension=".csv", filetypes=[("CSV文件", "*.csv"), ("所有文件", "*.*")])

    if not export_path:
        return

    try:
        with open(export_path, mode='w', newline='', encoding='utf-8-sig') as file:
            csv_writer = csv.writer(file)

            # 格式化时间
            formatted_start_time = time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(start_time))
            total_time = round(time.time() - start_time, 2)

            # 写入行标题和对应的值
            csv_writer.writerow(["属性", "值"])  # 标题行
            csv_writer.writerow(["目标图片路径", target_image_path])
            csv_writer.writerow(["选择的搜索目录", selected_directory.get()])
            csv_writer.writerow(["搜索开始时间", formatted_start_time])
            csv_writer.writerow(["所用时间(秒)", total_time])
            csv_writer.writerow(["支持的图片总数", len(processed_files)])
            csv_writer.writerow(["匹配图片总数", len(matched_images_list)])
            csv_writer.writerow([])  # 空行
            csv_writer.writerow(["匹配图片路径", "相似度"])  # 写入匹配图片路径和相似度列

            # 写入每个匹配图片的路径和相似度
            for matched_image, similarity in matched_images_with_similarity_list:
                csv_writer.writerow([matched_image, similarity])

        messagebox.showinfo("成功", f"匹配结果已导出到 {export_path}")
        logging.info(f"CSV匹配结果已导出至 {export_path}, 文件大小: {os.path.getsize(export_path) // 1024} KB")

    except Exception as e:
        messagebox.showerror("错误", f"导出失败: {e}")
        logging.error(f"导出失败: {e}")


# 导出为 Excel 文件并居中对齐
def export_to_excel():
    if not matched_images_list:
        messagebox.showinfo("提示", "没有可以导出的匹配结果。")
        return

    export_path = filedialog.asksaveasfilename(defaultextension=".xlsx", filetypes=[("Excel文件", "*.xlsx"), ("所有文件", "*.*")])

    if not export_path:
        return

    try:
        # 创建一个新的工作簿和工作表
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "匹配结果"

        # 格式化时间
        formatted_start_time = time.strftime('%Y/%m/%d %H:%M:%S', time.localtime(start_time))
        total_time = round(time.time() - start_time, 2)

        # 写入标题和数据
        data = [
            ["属性", "值"],
            ["目标图片路径", target_image_path],
            ["选择的搜索目录", selected_directory.get()],
            ["搜索开始时间", formatted_start_time],
            ["所用时间(秒)", total_time],
            ["支持的图片总数", len(processed_files)],
            ["匹配图片总数", len(matched_images_list)],
            [],
            ["匹配图片路径", "相似度"]
        ]

        # 先写入静态数据
        for row in data:
            ws.append(row)

        # 写入每个匹配图片的路径和相似度
        for matched_image, similarity in matched_images_with_similarity_list:
            ws.append([matched_image, similarity])

        # 设置单元格居中对齐
        alignment = Alignment(horizontal='center', vertical='center')

        # 对整个表格应用居中对齐
        for row in ws.iter_rows(min_row=1, max_row=ws.max_row, min_col=1, max_col=2):
            for cell in row:
                cell.alignment = alignment

        # 自动调整列宽
        for column_cells in ws.columns:
            # 获取该列中所有单元格内容的长度
            max_length = 0
            column = column_cells[0].column_letter  # 获取列字母
            for cell in column_cells:
                try:
                    if cell.value:
                        # 计算单元格内容的长度
                        cell_length = len(str(cell.value))
                        if cell_length > max_length:
                            max_length = cell_length
                except:
                    pass
            # 设置列宽，增加2个单位的余量
            adjusted_width = (max_length + 2)
            ws.column_dimensions[column].width = adjusted_width

        # 保存 Excel 文件
        wb.save(export_path)
        logging.info(f"Excel匹配结果已导出至 {export_path}, 文件大小: {os.path.getsize(export_path) // 1024} KB")
        messagebox.showinfo("成功", f"匹配结果已导出到 {export_path}")

    except Exception as e:
        logging.error(f"导出失败: {e}")
        messagebox.showerror("错误", f"导出失败: {e}")

# 显示选中的匹配图片
def show_selected_image(event):
    selected_index = result_list.curselection()
    if selected_index:
        selected_image_path = matched_images_list[selected_index[0]]  # 使用列表
        matched_img_label.config(text="加载中...")
        matched_img_label.update()  # 更新显示
        img = Image.open(selected_image_path)
        img.thumbnail((400, 400), Image.ANTIALIAS)  # 保持原比例缩放
        img = ImageTk.PhotoImage(img)
        matched_img_label.config(image=img)
        matched_img_label.image = img

# 禁用按钮
def disable_buttons():
    choose_target_btn.config(state=tk.DISABLED)
    choose_dir_btn.config(state=tk.DISABLED)
    start_btn.config(state=tk.DISABLED)
    stop_btn.config(state=tk.NORMAL)
    continue_btn.config(state=tk.DISABLED)  # 禁用继续按钮

# 启用按钮
def enable_buttons():
    logging.info("用户点击开始搜索按钮")
    choose_target_btn.config(state=tk.NORMAL)
    choose_dir_btn.config(state=tk.NORMAL)
    start_btn.config(state=tk.NORMAL)
    stop_btn.config(state=tk.DISABLED)
    continue_btn.config(state=tk.NORMAL)  # 启用继续按钮

# 停止搜索（单线程中不需要）
def stop_search():
    global is_searching
    logging.info("用户点击停止搜索按钮")
    is_searching = False
    progress_var.set("搜索已停止。")
    stop_event.set()  # 设置停止事件
    enable_buttons()

# 继续搜索
def continue_search(search_directory):
    global is_searching
    logging.info("用户点击继续搜索按钮")
    stop_event.clear()  # 清除停止事件
    is_searching = True  # 设置为继续搜索状态

    # 启动线程继续搜索
    search_thread = threading.Thread(target=search_for_matches, args=(search_directory,))
    search_thread.start()

# 定义打开文件夹的函数
def open_directory(path):
    if os.path.exists(path):
        os.startfile(path)  # 在 Windows 上打开文件资源管理器
    else:
        messagebox.showerror("错误", "路径不存在或已删除。")

# GUI 布局
frame = tk.Frame(root)
frame.pack(pady=20, expand=True, fill=tk.BOTH)

# 左侧部分
left_frame = tk.Frame(frame)
left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

# 选择目标图片和检索文件夹
bottom_frame = tk.Frame(left_frame)
bottom_frame.pack(side=tk.BOTTOM, fill=tk.X)

choose_target_btn = tk.Button(bottom_frame, text="选择目标图片", command=choose_target_image)
choose_target_btn.pack(side=tk.LEFT, padx=10)

choose_dir_btn = tk.Button(bottom_frame, text="选择搜索目录", command=choose_search_directory)
choose_dir_btn.pack(side=tk.LEFT, padx=10)

# 第一行：在选择搜索目录的按钮右边显示当前选择的文件夹
directory_label = tk.Label(bottom_frame, textvariable=selected_directory, width=50, anchor="w")
directory_label.pack(side=tk.TOP, fill=tk.X, padx=10, pady=2)  # 使用 pady 增加一点间距

# 第二行：实时显示当前正在扫描的子目录路径
current_directory_label = tk.Label(bottom_frame, textvariable=current_directory, width=50, anchor="w", fg="blue")
current_directory_label.pack(side=tk.TOP, fill=tk.X, padx=10, pady=2)

# 绑定鼠标点击事件
directory_label.bind("<Button-1>", lambda e: open_directory(selected_directory.get()))
current_directory_label.bind("<Button-1>", lambda e: open_directory(current_directory.get()))

# 显示选中的目标图片
target_img_label = tk.Label(left_frame, text="目标图片", width=30, height=15, relief=tk.RIDGE)
target_img_label.pack(pady=10, fill=tk.BOTH, expand=True)

# 在左侧部分添加滚动条的容器
listbox_frame = tk.Frame(left_frame)
listbox_frame.pack(pady=10, fill=tk.BOTH, expand=True)

# 创建匹配结果的列表框
result_list = tk.Listbox(listbox_frame, xscrollcommand=lambda f, l: x_scrollbar.set(f, l), yscrollcommand=lambda f, l: y_scrollbar.set(f, l))
result_list.grid(row=0, column=0, sticky="nsew")  # 使用 grid() 管理器放置 Listbox
result_list.bind('<<ListboxSelect>>', show_selected_image)

# 创建垂直滚动条
y_scrollbar = tk.Scrollbar(listbox_frame, orient=tk.VERTICAL, command=result_list.yview)
y_scrollbar.grid(row=0, column=1, sticky="ns")  # 将垂直滚动条放置在 Listbox 旁边

# 创建水平滚动条
x_scrollbar = tk.Scrollbar(listbox_frame, orient=tk.HORIZONTAL, command=result_list.xview)
x_scrollbar.grid(row=1, column=0, columnspan=2, sticky="ew")  # 将水平滚动条放置在 Listbox 底部

# 将滚动条与列表框关联
result_list.config(yscrollcommand=y_scrollbar.set, xscrollcommand=x_scrollbar.set)

# 设置 grid 的行列权重，使得 Listbox 可以扩展
listbox_frame.grid_rowconfigure(0, weight=1)  # 使 Listbox 行可扩展
listbox_frame.grid_columnconfigure(0, weight=1)  # 使 Listbox 列可扩展

# 右侧部分
right_frame = tk.Frame(frame)
right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)

# 显示匹配图片
matched_img_label = tk.Label(right_frame, text="匹配图片预览", width=30, height=15, relief=tk.RIDGE)
matched_img_label.pack(pady=10, fill=tk.BOTH, expand=True)

# 进度显示
progress_var = tk.StringVar()
progress_label = tk.Label(right_frame, textvariable=progress_var)
progress_label.pack(pady=10)

# 进度条
progress_bar = Progressbar(right_frame, orient="horizontal", length=300, mode="determinate")
progress_bar.pack(pady=10)

# 按钮布局，位于右下角
button_frame = tk.Frame(right_frame)
button_frame.pack(side=tk.BOTTOM, fill=tk.X, padx=10, pady=10)

# 停止搜索按钮
stop_btn = tk.Button(button_frame, text="结束搜索", fg="black", bg="red", command=stop_search, state=tk.DISABLED)
stop_btn.pack(side=tk.RIGHT, padx=10)

# 继续搜索按钮
continue_btn = tk.Button(button_frame, text="继续搜索", fg="white", bg="blue", command=lambda: continue_search(selected_directory.get()), state=tk.DISABLED)
continue_btn.pack(side=tk.RIGHT, padx=10)

# 开始搜索按钮
start_btn = tk.Button(button_frame, text="开始搜索", fg="white", bg="green", command=choose_search_directory)
start_btn.pack(side=tk.RIGHT, padx=10)

# 添加导出按钮
export_btn = tk.Button(button_frame, text="导出CSV", fg="black", bg="orange", command=export_to_csv)
export_btn.pack(side=tk.LEFT, padx=10)

# 添加导出按钮
export_btn = tk.Button(button_frame, text="导出Excel", fg="black", bg="orange", command=export_to_excel)
export_btn.pack(side=tk.LEFT, padx=10)

# 添加“导出所有匹配图片”按钮
export_images_btn = tk.Button(button_frame, text="导出所有匹配图片", fg="white", bg="orange", command=export_matched_images)
export_images_btn.pack(side=tk.LEFT, padx=10)

# 将查看日志的按钮添加到按钮区域
log_btn = tk.Button(button_frame, text="查看日志", fg="white", bg="purple", command=open_log_file)
log_btn.pack(side=tk.RIGHT, padx=10)

# 初始化UI更新队列
root.after(100, update_ui_from_queue)


root.mainloop()
